import csv
import ast
import os
import time
import pickle
from itertools import combinations
from collections import defaultdict
import functools


class Apriori:
    def __init__(self, min_support=0.1, min_confidence=0.7):
        self.min_support = min_support
        self.min_confidence = min_confidence
        self.freq_itemsets = []
        self.rules = []
        self.total_transactions = 0
        self.transactions = []

    @classmethod
    def load(cls, filename):
        """增强加载方法的健壮性"""
        try:
            with open(filename, 'rb') as f:
                data = pickle.load(f)
            instance = cls()
            instance.freq_itemsets = data['freq_itemsets']
            instance.rules = data['rules']
            instance.total_transactions = data['total_transactions']
            return instance
        except Exception as e:
            print(f"加载失败: {str(e)}")
            return cls()

    def save(self, filename='apriori.bin'):
        """新增保存方法"""
        data = {
            'freq_itemsets': self.freq_itemsets,
            'rules': self.rules,
            'total_transactions': self.total_transactions
        }
        with open(filename, 'wb') as f:
            pickle.dump(data, f)

    def _load_data(self):
        """增强数据加载的容错性"""
        data = []
        try:
            with open('processed_招聘数据.csv', 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    try:
                        skills = ast.literal_eval(row['skill_list'])
                        cleaned = [s.strip().lower() for s in skills if s.strip()]
                        if cleaned:
                            data.append(frozenset(cleaned))
                    except (SyntaxError, ValueError) as e:
                        print(f"数据解析错误: {str(e)}")
                        continue
        except FileNotFoundError:
            print("数据文件未找到")
        return data

    def _generate_candidates(self, itemset, length):
        """优化候选生成算法"""
        return {i.union(j) for i in itemset for j in itemset
                if len(i.union(j)) == length and i != j}

    def _prune(self, itemset, prev_itemsets):
        """增强剪枝效率"""
        return {item for item in itemset
                if all(frozenset(comb) in prev_itemsets
                       for comb in combinations(item, len(item) - 1))}

    def _get_freq_itemsets(self, transactions):
        """优化Apriori算法实现"""
        freq_items = []
        k = 1

        # 生成1-项集
        item_counts = defaultdict(int)
        for transaction in transactions:
            for item in transaction:
                item_counts[frozenset([item])] += 1

        current_items = {item: count for item, count in item_counts.items()
                         if count / self.total_transactions >= self.min_support}
        freq_items.append(current_items)

        # 生成高阶项集
        while current_items:
            k += 1
            candidates = self._generate_candidates(current_items.keys(), k)
            candidates = self._prune(candidates, current_items.keys())

            counts = defaultdict(int)
            for transaction in transactions:
                for candidate in candidates:
                    if candidate.issubset(transaction):
                        counts[candidate] += 1

            current_items = {k: v for k, v in counts.items()
                             if v / self.total_transactions >= self.min_support}
            if current_items:
                freq_items.append(current_items)

        return freq_items

    def _generate_rules(self):
        """优化规则生成算法"""
        rules = []
        for level in self.freq_itemsets[1:]:
            for itemset, support in level.items():
                subsets = (frozenset(s) for l in range(1, len(itemset))
                           for s in combinations(itemset, l))

                for antecedent in subsets:
                    consequent = itemset - antecedent
                    if not consequent:
                        continue

                    antecedent_support = self._get_support(antecedent)
                    if antecedent_support == 0:
                        continue

                    confidence = support / antecedent_support
                    consequent_support = self._get_support(consequent)
                    lift = confidence / (consequent_support / self.total_transactions) if consequent_support else 0

                    if confidence >= self.min_confidence:
                        rules.append({
                            'antecedent': antecedent,
                            'consequent': consequent,
                            'support': support / self.total_transactions,
                            'confidence': confidence,
                            'lift': lift
                        })
        return rules

    def _get_support(self, itemset):
        """优化支持度查询"""
        for level in reversed(self.freq_itemsets):
            for candidate, count in level.items():
                if candidate == itemset:
                    return count
        return 0

    def run(self):
        """增强运行方法的可靠性"""
        if os.path.exists('apriori.bin'):
            try:
                cached = self.load('apriori.bin')
                self.__dict__.update(cached.__dict__)
                print("成功加载缓存")
                return
            except Exception as e:
                print(f"缓存加载失败: {str(e)}")

        print("开始重新计算...")
        start_time = time.time()

        self.transactions = self._load_data()
        self.total_transactions = len(self.transactions)
        print(f"加载{self.total_transactions}条事务记录")

        if self.total_transactions == 0:
            raise ValueError("无有效数据可供分析")

        self.freq_itemsets = self._get_freq_itemsets(self.transactions)
        self.rules = self._generate_rules()
        self.save()

        print(f"计算完成，耗时{time.time() - start_time:.2f}秒")

    def get_associated_skills(self, keyword, min_support=0.1):
        """优化关联查询方法"""
        keyword = keyword.strip().lower()
        if not self.freq_itemsets:
            raise ValueError("请先运行分析过程")

        associations = defaultdict(float)

        for level in self.freq_itemsets:
            for itemset, count in level.items():
                if keyword in itemset and len(itemset) >= 2:
                    support = count / self.total_transactions
                    for skill in itemset:
                        if skill != keyword and support >= min_support:
                            associations[skill] = max(associations[skill], support)

        # 修改排序方式
        sorted_items = sorted(associations.items(),
                              key=lambda x: (-x[1], x[0]))

        # 格式化支持度字符串
        return {skill: f"{support:.2%}"
                for skill, support in sorted_items}

    def show_results(self):
        """优化结果展示"""
        print("\n=== 频繁项集分析 ===")
        for level, itemsets in enumerate(self.freq_itemsets, 1):
            print(f"\n{level}-项集 (共{len(itemsets)}个):")
            for itemset, count in itemsets.items():
                print(f"  {set(itemset)}: 支持度 {count / self.total_transactions:.2%}")

        print("\n=== 前20强关联规则 ===")
        sorted_rules = sorted(self.rules,
                              key=lambda x: (-x['confidence'], -x['support']))
        for i, rule in enumerate(sorted_rules[:20], 1):
            print(f"\n规则 #{i}:")
            print(f"{set(rule['antecedent'])} => {set(rule['consequent'])}")
            print(f"置信度: {rule['confidence']:.2%}  支持度: {rule['support']:.2%}")
            print(f"提升度: {rule['lift']:.2f}")


def load_apriori(filename='apriori.bin'):
    """增强加载函数可靠性"""
    return Apriori.load(filename)


def getAsso(apriori_instance, keyword, min_support=0.1):
    """优化关联查询接口"""
    raw = apriori_instance.get_associated_skills(keyword, min_support)
    return {f"{k}（支持度 {v}）": v for k, v in raw.items()}


if __name__ == "__main__":
    analyzer = Apriori()
    analyzer.run()

    # 示例分析
    analyzer.show_results()

    # 交互查询
    while True:
        query = input("\n输入查询技能（q退出）: ").strip()
        if query.lower() == 'q':
            break
        if result := analyzer.get_associated_skills(query, 0.1):
            print(f"\n与 {query} 关联的技能：")
            for skill, support in result.items():
                print(f"  {skill}: {support}")
        else:
            print("未找到相关关联")